# memory_optimization.py
import torch
from typing import Dict, List


class MemoryOptimizer:
    """内存优化器"""

    def __init__(self):
        self.gpu_memory = torch.cuda.get_device_properties(0).total_memory
        self.used_memory = 0

    def estimate_memory_usage(self, model_config, batch_size: int) -> int:
        """估算内存使用"""
        # 模型权重内存
        model_params = sum(p.numel() for p in model_config.model.parameters())
        model_memory = model_params * 2  # FP16

        # KV缓存内存
        kv_cache_memory = self._estimate_kv_cache_memory(model_config, batch_size)

        # 激活内存
        activation_memory = self._estimate_activation_memory(model_config, batch_size)

        return model_memory + kv_cache_memory + activation_memory

    def _estimate_kv_cache_memory(self, model_config, batch_size: int) -> int:
        """估算KV缓存内存"""
        num_layers = model_config.config.num_hidden_layers
        num_heads = model_config.config.num_attention_heads
        head_dim = model_config.config.hidden_size // num_heads
        seq_len = model_config.max_seq_len

        # 每层的KV缓存大小
        layer_kv_memory = 2 * batch_size * seq_len * num_heads * head_dim * 2  # bytes

        return num_layers * layer_kv_memory

    def optimize_batch_size(self, model_config) -> int:
        """优化批次大小"""
        available_memory = self.gpu_memory - self.used_memory

        # 二分查找最大批次大小
        low, high = 1, 32
        best_batch_size = 1

        while low <= high:
            mid = (low + high) // 2
            memory_needed = self.estimate_memory_usage(model_config, mid)

            if memory_needed <= available_memory:
                best_batch_size = mid
                low = mid + 1
            else:
                high = mid - 1

        return best_batch_size